{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "64594732",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-21T23:24:45.869340900Z",
     "start_time": "2025-06-21T23:10:13.353784700Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-chinese and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "Be aware, overflowing tokens are not returned for the setting you have chosen, i.e. sequence pairs with the 'longest_first' truncation strategy. So the returned list will always be empty even if some tokens have been removed.\n",
      "Be aware, overflowing tokens are not returned for the setting you have chosen, i.e. sequence pairs with the 'longest_first' truncation strategy. So the returned list will always be empty even if some tokens have been removed.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ 预测完成，结果保存在 submission.csv\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import torch\n",
    "from transformers import BertTokenizer, BertForSequenceClassification\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import pandas as pd\n",
    "\n",
    "# 配置\n",
    "model_path = \"./bert_model.pth\"  # 模型路径\n",
    "test_file = \"gaiic_track3_round1_testA_20210228.tsv\"  # 测试数据路径\n",
    "output_file = \"submission.csv\"  # 输出文件路径\n",
    "max_len = 64  # 输入最大长度\n",
    "batch_size = 32\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# 数据集定义\n",
    "class TestTextDataset(Dataset):\n",
    "    def __init__(self, file_path, tokenizer, max_len):\n",
    "        self.data = pd.read_csv(file_path, sep=\"\\t\", header=None, names=[\"query1\", \"query2\"])\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_len = max_len\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        row = self.data.iloc[idx]\n",
    "        inputs = self.tokenizer(\n",
    "            row[\"query1\"],\n",
    "            row[\"query2\"],\n",
    "            max_length=self.max_len,\n",
    "            padding=\"max_length\",\n",
    "            truncation=True,\n",
    "            return_tensors=\"pt\"\n",
    "        )\n",
    "        return {\n",
    "            'input_ids': inputs['input_ids'].squeeze(0),\n",
    "            'attention_mask': inputs['attention_mask'].squeeze(0)\n",
    "        }\n",
    "\n",
    "# 初始化模型与tokenizer\n",
    "tokenizer = BertTokenizer.from_pretrained(\"bert-base-chinese\")\n",
    "model = BertForSequenceClassification.from_pretrained(\"bert-base-chinese\", num_labels=2)\n",
    "model.load_state_dict(torch.load(model_path, map_location=device))\n",
    "model.to(device)\n",
    "model.eval()\n",
    "\n",
    "# 加载测试集\n",
    "test_dataset = TestTextDataset(test_file, tokenizer, max_len)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "# 预测\n",
    "predictions = []\n",
    "with torch.no_grad():\n",
    "    for batch in test_loader:\n",
    "        input_ids = batch['input_ids'].to(device)\n",
    "        attention_mask = batch['attention_mask'].to(device)\n",
    "\n",
    "        outputs = model(input_ids=input_ids, attention_mask=attention_mask)\n",
    "        probs = torch.softmax(outputs.logits, dim=1)\n",
    "        predictions.extend(probs[:, 1].cpu().numpy())\n",
    "\n",
    "# 保存预测结果\n",
    "pd.DataFrame(predictions).to_csv(output_file, index=False, header=False)\n",
    "print(\"✅ 预测完成，结果保存在 submission.csv\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "language": "python",
   "display_name": "Python 3 (ipykernel)"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
